#pragma once

#include "nanoflann.hpp"

using namespace nanoflann;

#include "geometry/Vector.h"

typedef std::pair<size_t, double> KDResultPair;
typedef std::vector< KDResultPair  > KDResults;

template <typename T, int N>
class NanoflannKdTree{

public:
	typedef Vector<T, N> VectorN;

	struct PointCloud{
		std::vector<VectorN>  pts;

		inline size_t kdtree_get_point_count() const { return pts.size(); }

		// using squared distance
		inline T kdtree_distance(const T *p1, const size_t idx_p2, size_t size) const{
			VectorN p2 = pts[idx_p2];
			T d(0), di;
			for (size_t i=0; i<size; i++)
			{
				di = p1[i] - p2[i];
				d += di * di;
			}
			return d;
		}

		inline T kdtree_get_pt(const size_t idx, int dim) const{
			return pts[idx][dim];
		}

		template <class BBOX>
		bool kdtree_get_bbox(BBOX &) const { return false; }
	};

	PointCloud cloud;

public:
	typedef KDTreeSingleIndexAdaptor< L2_Simple_Adaptor<T, PointCloud>, PointCloud, N /* dim */ > my_kd_tree;
	my_kd_tree * tree;

    NanoflannKdTree(){
        tree = NULL;
    }

    ~NanoflannKdTree(){
        delete tree;
    }

	void addPoint(const VectorN & p){
		cloud.pts.push_back(p);
	}

	void build()
	{
        if(tree) delete tree;

		// construct a kd-tree index:
		tree = new my_kd_tree(N /*dim*/, cloud, KDTreeSingleIndexAdaptorParams(10 /* max leaf */) );
		tree->buildIndex();
	}

	size_t k_closest(VectorN p, int k, KDResults & ret_matches)
	{
        k = k < (int)cloud.pts.size() ? k : cloud.pts.size();

		ret_matches.clear();
		ret_matches.resize(k);

		std::vector<size_t> ret_index(k);
		std::vector<T> out_dist(k);

		tree->knnSearch(&p[0], k, &ret_index[0], &out_dist[0]);

		// id + dis
		for(int i = 0; i < k; i++)
			ret_matches[i] = std::make_pair(ret_index[i], sqrt(out_dist[i]));

		return k;
	}

    size_t ball_search(VectorN p, T search_radius)
    {
        KDResults ret_matches;
        return ball_search(p, search_radius, ret_matches);
    }

    size_t ball_search(VectorN p, T search_radius, KDResults & ret_matches)
	{
		ret_matches.clear();

		nanoflann::SearchParams params;

		// radius is squared since the use of squared distance
		tree->radiusSearch(&p[0], search_radius * search_radius, ret_matches, params);

		// get the correct dis
		for(size_t i = 0; i < ret_matches.size(); i++)
			ret_matches[i].second = sqrt(ret_matches[i].second);

		return ret_matches.size();
	}

    int closest(VectorN & p)
    {
        KDResults match;
        this->k_closest(p, 1, match);

        if(!match.size())
            return -1;
        else
            return match[0].first;
    }

	inline bool has(VectorN & p, T eps_distance = 1e-7)
	{
		KDResults match;
		ball_search(p, eps_distance, match);
		return match.size();
	}
};

typedef NanoflannKdTree<double, 2> NanoflannKdTree2d;
typedef NanoflannKdTree<double, 3> NanoflannKdTree3d;
